/**
 * Copyright (c) 2021 Darius Rückert
 * Licensed under the MIT License.
 * See LICENSE file for more information.
 */

#pragma once
#include "saiga/core/math/math.h"
/**
 * Functions for finding the real roots of polynomials up to degree 4.
 *
 * Not exactly sure where the code below is from.
 * Opencv also includes the same exact code without specifying the original author.
 * See here:
 * opencv/modules/calib3d/src/polynom_solver.h
 */
namespace Saiga
{
inline int solve_deg2(double a, double b, double c, double& x1, double& x2)
{
    double delta = b * b - 4 * a * c;

    if (delta < 0) return 0;

    double inv_2a = 0.5 / a;

    if (delta == 0)
    {
        x1 = -b * inv_2a;
        x2 = x1;
        return 1;
    }

    double sqrt_delta = sqrt(delta);
    x1                = (-b + sqrt_delta) * inv_2a;
    x2                = (-b - sqrt_delta) * inv_2a;
    return 2;
}


/// Reference : Eric W. Weisstein. "Cubic Equation." From MathWorld--A Wolfram Web Resource.
/// http://mathworld.wolfram.com/CubicEquation.html
/// \return Number of real roots found.
inline int solve_deg3(double a, double b, double c, double d, double& x0, double& x1, double& x2)
{
    if (a == 0)
    {
        // Solve second order sytem
        if (b == 0)
        {
            // Solve first order system
            if (c == 0) return 0;

            x0 = -d / c;
            return 1;
        }

        x2 = 0;
        return solve_deg2(b, c, d, x0, x1);
    }

    // Calculate the normalized form x^3 + a2 * x^2 + a1 * x + a0 = 0
    double inv_a = 1. / a;
    double b_a = inv_a * b, b_a2 = b_a * b_a;
    double c_a = inv_a * c;
    double d_a = inv_a * d;

    // Solve the cubic equation
    double Q     = (3 * c_a - b_a2) / 9;
    double R     = (9 * b_a * c_a - 27 * d_a - 2 * b_a * b_a2) / 54;
    double Q3    = Q * Q * Q;
    double D     = Q3 + R * R;
    double b_a_3 = (1. / 3.) * b_a;

    if (Q == 0)
    {
        if (R == 0)
        {
            x0 = x1 = x2 = -b_a_3;
            return 3;
        }
        else
        {
            x0 = pow(2 * R, 1 / 3.0) - b_a_3;
            return 1;
        }
    }

    if (D <= 0)
    {
        // Three real roots
        double theta  = acos(R / sqrt(-Q3));
        double sqrt_Q = sqrt(-Q);
        x0            = 2 * sqrt_Q * cos(theta / 3.0) - b_a_3;
        x1            = 2 * sqrt_Q * cos((theta + 2 * pi<double>()) / 3.0) - b_a_3;
        x2            = 2 * sqrt_Q * cos((theta + 4 * pi<double>()) / 3.0) - b_a_3;

        return 3;
    }

    // D > 0, only one real root
    double AD = pow(fabs(R) + sqrt(D), 1.0 / 3.0) * (R > 0 ? 1 : (R < 0 ? -1 : 0));
    double BD = (AD == 0) ? 0 : -Q / AD;

    // Calculate the only real root
    x0 = AD + BD - b_a_3;

    return 1;
}

/// Reference : Eric W. Weisstein. "Quartic Equation." From MathWorld--A Wolfram Web Resource.
/// http://mathworld.wolfram.com/QuarticEquation.html
/// \return Number of real roots found.
inline int solve_deg4(double a, double b, double c, double d, double e, double& x0, double& x1, double& x2, double& x3)
{
    if (a == 0)
    {
        x3 = 0;
        return solve_deg3(b, c, d, e, x0, x1, x2);
    }

    // Normalize coefficients
    double inv_a = 1. / a;
    b *= inv_a;
    c *= inv_a;
    d *= inv_a;
    e *= inv_a;
    double b2 = b * b, bc = b * c, b3 = b2 * b;

    // Solve resultant cubic
    double r0, r1, r2;
    int n = solve_deg3(1, -c, d * b - 4 * e, 4 * c * e - d * d - b2 * e, r0, r1, r2);
    if (n == 0) return 0;

    // Calculate R^2
    double R2 = 0.25 * b2 - c + r0, R;
    if (R2 < 0) return 0;

    R            = sqrt(R2);
    double inv_R = 1. / R;

    int nb_real_roots = 0;

    // Calculate D^2 and E^2
    double D2, E2;
    if (R < 10E-12)
    {
        double temp = r0 * r0 - 4 * e;
        if (temp < 0)
            D2 = E2 = -1;
        else
        {
            double sqrt_temp = sqrt(temp);
            D2               = 0.75 * b2 - 2 * c + 2 * sqrt_temp;
            E2               = D2 - 4 * sqrt_temp;
        }
    }
    else
    {
        double u = 0.75 * b2 - 2 * c - R2, v = 0.25 * inv_R * (4 * bc - 8 * d - b3);
        D2 = u + v;
        E2 = u - v;
    }

    double b_4 = 0.25 * b, R_2 = 0.5 * R;
    if (D2 >= 0)
    {
        double D      = sqrt(D2);
        nb_real_roots = 2;
        double D_2    = 0.5 * D;
        x0            = R_2 + D_2 - b_4;
        x1            = x0 - D;
    }

    // Calculate E^2
    if (E2 >= 0)
    {
        double E   = sqrt(E2);
        double E_2 = 0.5 * E;
        if (nb_real_roots == 0)
        {
            x0            = -R_2 + E_2 - b_4;
            x1            = x0 - E;
            nb_real_roots = 2;
        }
        else
        {
            x2            = -R_2 + E_2 - b_4;
            x3            = x2 - E;
            nb_real_roots = 4;
        }
    }

    return nb_real_roots;
}

inline int solve_deg4(const Eigen::Matrix<double, 5, 1>& coeffs, Vec4& roots_real)
{
    return solve_deg4(coeffs(0), coeffs(1), coeffs(2), coeffs(3), coeffs(4), roots_real(0), roots_real(1),
                      roots_real(2), roots_real(3));
}

}  // namespace Saiga
